#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <assert.h>
#include <mpi.h>

typedef struct sSource {
    int row;
    int col;
    double amount;
} Source;

Source* get_input(int my_rank,
                  int* global_n_p,
                  int* num_steps_p,
                  int* num_srcs_p);

void set_output(int my_rank,
                int comm_sz,
                int local_n,
                int global_n,
                double** my_rows);

void reset_sources(int my_rank,
                   int local_n,
                   int global_n,
                   int num_srcs,
                   Source* srcs,
                   double** my_rows);

void update_cell(int my_rank,
                 int comm_sz,
                 int local_n,
                 int global_n,
                 int row,
                 int col,
                 double** my_rows,
                 double** tmp_rows,
                 double* sup_row,
                 double* inf_row);

double** create2DArray(int row_size, int col_size);
void remove2DArray(int row_size, int col_size, double** array);


int main(int argc, char **argv)
{
    int  comm_sz, my_rank;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &comm_sz);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    int global_n, num_steps, num_srcs;
    Source* srcs;
    srcs = get_input(my_rank, &global_n, &num_steps, &num_srcs);

    int local_n;
    local_n = global_n/comm_sz; /*  number of rows */

    double** my_rows;  /* local matrix */
    my_rows = create2DArray(local_n, global_n);

    // Set the sources for the first time
    reset_sources(my_rank, local_n, global_n, num_srcs, srcs, my_rows);

    double** swapper = NULL;  /* pointer for swapping */
    double** tmp_rows;        /* temporal matrix */
    tmp_rows = create2DArray(local_n, global_n);

    double* sup_row;
    double* inf_row;
    // sup_row and inf_row are set to 0
    sup_row = (double*) calloc(global_n, sizeof(double));
    inf_row = (double*) calloc(global_n, sizeof(double));

    int step, row, col;
    for (step=0; step<num_steps; step++) {

        if (comm_sz > 1){
            // Exchange the inferior row
            if (my_rank < comm_sz-1) {
                MPI_Sendrecv(my_rows[local_n-1], global_n, MPI_DOUBLE,
                             my_rank+1, 0, inf_row, global_n, MPI_DOUBLE,
                             my_rank+1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            }
            // // Exchange the superior row
            if (my_rank > 0) {
                MPI_Sendrecv(my_rows[0], global_n, MPI_DOUBLE,
                             my_rank-1, 0, sup_row, global_n, MPI_DOUBLE,
                             my_rank-1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            }
        }

        // Update my_rows, stores in tmp_rows
        for (row=0; row<local_n; row++) {
            for (col=0; col<global_n; col++) {
                update_cell(my_rank, comm_sz, local_n, global_n, row, col,
                            my_rows, tmp_rows, sup_row, inf_row);
            }
        }

        // Refresh the soruces
        reset_sources(my_rank, local_n, global_n, num_srcs, srcs, tmp_rows);

        // Swap between tmp_rows and my_rows
        swapper = my_rows;
        my_rows = tmp_rows;
        tmp_rows = swapper;
    }

    // Write output to a file
    set_output(my_rank, comm_sz, local_n, global_n, my_rows);

    free(srcs);
    free(sup_row);
    free(inf_row);
    remove2DArray(local_n, global_n, my_rows);
    remove2DArray(local_n, global_n, tmp_rows);

    MPI_Finalize();
    return 0;
}


Source* get_input(int my_rank,
                  int* global_n_p,
                  int* num_steps_p,
                  int* num_srcs_p)
{
    /* Parses input from user, and distrubutes to all processes */

    Source* srcs;

    if (my_rank == 0) {
        int n;
        printf("Enter global_n, num_steps and num_srcs\n");
        n = scanf("%d", global_n_p);
        assert(n == 1);
        n = scanf("%d", num_steps_p);
        assert(n == 1);
        n = scanf("%d", num_srcs_p);
        assert(n == 1);

        int num_srcs = *num_srcs_p;
        srcs = (Source*) malloc(num_srcs * sizeof(Source));

        int i, x, y;
        double amt;
        for (i=0; i<num_srcs; i++){
            printf("Enter the triple (x, y, t) number %d of %d\n", i, num_srcs);
            n = scanf("%d %d %lf", &x, &y, &amt);
            assert(n == 3);

            srcs[i].row = x;
            srcs[i].col = y;
            srcs[i].amount = amt;
        }
        printf("\n");
    }

    // Create a mpi type for the sources
    int blocklengths[3] = {1, 1, 1};
    MPI_Datatype types[3] = {MPI_INT, MPI_INT, MPI_DOUBLE};
    MPI_Datatype mpi_source;
    MPI_Aint offsets[3];

    offsets[0] = offsetof(Source, row);
    offsets[1] = offsetof(Source, col);
    offsets[2] = offsetof(Source, amount);

    MPI_Type_create_struct(3, blocklengths, offsets, types, &mpi_source);
    MPI_Type_commit(&mpi_source);
    // End

    MPI_Bcast(global_n_p, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(num_steps_p, 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(num_srcs_p, 1, MPI_INT, 0, MPI_COMM_WORLD);

    if (my_rank != 0) {
        int num_srcs = *num_srcs_p;
        srcs = (Source*) malloc(num_srcs * sizeof(Source));
    }

    MPI_Bcast(srcs, *num_srcs_p, mpi_source, 0, MPI_COMM_WORLD);

    MPI_Type_free(&mpi_source);
    return srcs;
}

void set_output(int my_rank,
                int comm_sz,
                int local_n,
                int global_n,
                double** my_rows)
{
    /* Obtains the local rows from all processes and writes them to a file */

    size_t file_length;
    FILE* output_file = fopen("output.csv", "w");
    double* row_to_print;
    row_to_print = (double*) calloc(global_n, sizeof(double));

    int proc, i, j, k;
    if (my_rank != 0) {
        for (i=0; i<local_n; i++) {
            MPI_Send(my_rows[i], global_n, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
        }

    } else { /* my_rank == 0 */
        /* First, write my rows */
        for (i=0; i<local_n; i++){
            for (j=0; j<global_n; j++) {
                    fprintf(output_file,"%f ", my_rows[i][j]);
            }
            fputs("\n", output_file);
        }

        /* For each process (except myself) */
        for (proc=1; proc<comm_sz; proc++) {
            /* For each row sent*/
            for (i=0; i<local_n; i++) {
                MPI_Recv(row_to_print, global_n, MPI_DOUBLE, proc, 0,
                        MPI_COMM_WORLD, MPI_STATUS_IGNORE);

                /* Write the row on file*/
                for (k=0; k<global_n; k++) {
                    fprintf(output_file,"%f ", row_to_print[k]);
                }
                fputs("\n", output_file);
            }
        }
    }
    fclose(output_file);
    free(row_to_print);
}


void reset_sources(int my_rank,
                   int local_n,
                   int global_n,
                   int num_srcs,
                   Source* srcs,
                   double** my_rows)
{
    /* Resets the source's ammount */

    // Map my local rows range to the global range
    int initial_pos = my_rank*local_n;
    int final_pos = initial_pos + local_n -1;

    int i, j, curr_row, curr_col;
    double curr_amount;
    // For all sources
    for (i=0; i<num_srcs; i++) {
        curr_row = srcs[i].row;
        curr_col = srcs[i].col;
        curr_amount = srcs[i].amount;

        // If curr_row belongs to my_rows
        if (initial_pos <= curr_row && curr_row <= final_pos
            && 0 <= curr_col && curr_col < global_n) {
            // Convert to a local row
            curr_row = curr_row - (my_rank*local_n);
            my_rows[curr_row][curr_col] = curr_amount;
        }
    }
}


void update_cell(int my_rank,
                 int comm_sz,
                 int local_n,
                 int global_n,
                 int row,
                 int col,
                 double** my_rows,
                 double** tmp_rows,
                 double* sup_row,
                 double* inf_row)
{
    /* Updates a cell from my_rows into tmp_rows */

    int divisor = 5;
    double up = 0;
    double down = 0;
    double left = 0;
    double right = 0;


    /* Upper bound checks */
    if (row > 0) {                        /* I'm not the first local */
        up = my_rows[row-1][col];
    } else {                              /*  I'm the first local */
        if (my_rank != 0 ){               /* I'm not the first global */
            up = sup_row[col];
        } else {                          /* I'm the first global */
            divisor--;
        }
    }

    /* Lower bound checks */
    if (row < local_n-1) {
        down = my_rows[row+1][col];
    } else {
        if(my_rank != comm_sz-1) {
            down = inf_row[col];
        } else {
            divisor--;
        }
    }

    /* Left bound checks */
    if (col > 0) {
        left = my_rows[row][col-1];
    } else {
        divisor--;
    }

    /* Right bound checks */
    if (col < global_n-1) {
        right = my_rows[row][col+1];
    } else {
        divisor--;
    }

    tmp_rows[row][col] = (my_rows[row][col]+up+down+left+right)/(double)divisor;
}


double** create2DArray(int row_size, int col_size)
{
    /* Creates an array of arrays */

    double** array;
    array = (double**) malloc(row_size * sizeof(double*));

    int i, j;
    for (i=0; i<row_size; i++) {
        // Set the matrix to 0
        array[i] = (double*) calloc(col_size, sizeof(double));
    }

    return array;
}


void remove2DArray(int row_size, int col_size, double** array)
{
    /* Destroys an array of arrays */

    int i;
    for (i=0; i<row_size; i++) {
       free(array[i]);
    }
    free(array);
}
